package io.focus.mijin.zuul.ratelimit.http;

import com.google.common.collect.Lists;
import com.google.common.collect.Maps;
import com.netflix.zuul.constants.ZuulHeaders;
import com.netflix.zuul.context.RequestContext;
import com.netflix.zuul.http.HttpServletRequestWrapper;
import com.netflix.zuul.http.ServletInputStreamWrapper;
import com.netflix.zuul.util.HTTPRequestUtils;
import org.apache.commons.lang.StringUtils;
import org.apache.tomcat.util.http.fileupload.IOUtils;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

import javax.servlet.ServletInputStream;
import javax.servlet.http.HttpServletRequest;
import java.io.*;
import java.net.SocketTimeoutException;
import java.net.URLDecoder;
import java.util.*;

/**
 * 重写zuul的HttpServletRequestWrapper
 *
 * @author 吴爱军
 * @date 2018/5/25
 */
public class RateLimitHttpServletRequestWrapper extends HttpServletRequestWrapper {

    private final static HashMap<String, String[]> EMPTY_MAP = Maps.newHashMap();

    protected static final Logger LOG = LoggerFactory.getLogger(RateLimitHttpServletRequestWrapper.class);

    private HttpServletRequest req;

    private byte[] contentData = null;

    private HashMap<String, String[]> parameters = null;

    private long bodyBufferingTimeNs = 0;

    public RateLimitHttpServletRequestWrapper() {
        super(groovyTrick());
    }

    private static HttpServletRequest groovyTrick() {
        //a trick for Groovy
        throw new IllegalArgumentException("Please use HttpServletRequestWrapper(HttpServletRequest request) constructor!");
    }

    private RateLimitHttpServletRequestWrapper(HttpServletRequest request, byte[] contentData, HashMap<String, String[]> parameters) {
        super(request);
        req = request;
        this.contentData = contentData;
        this.parameters = parameters;
    }

    public RateLimitHttpServletRequestWrapper(HttpServletRequest request) {
        super(request);
        req = request;
    }

    /**
     * Returns the wrapped HttpServletRequest.
     * Using the getParameterXXX(), getInputStream() or getReader() methods may interfere
     * with this class operation.
     *
     * @return The wrapped HttpServletRequest.
     */
    @Override
    public HttpServletRequest getRequest() {
        try {
            parseRequest();
        } catch (IOException e) {
            throw new IllegalStateException("Cannot parse the request!", e);
        }
        return req;
    }

    /**
     * This method is safe to use multiple times.
     *
     * @return The request body data.
     */
    @Override
    public byte[] getContentData() {
        return contentData;
    }

    /**
     * This method is safe to use multiple times.
     * Changing the returned map or the array of any of the map's values will not
     * interfere with this class operation.
     *
     * @return The cloned parameters map.
     */
    @Override
    public HashMap<String, String[]> getParameters() {
        if (parameters == null) {
            return EMPTY_MAP;
        }
        HashMap<String, String[]> map = new HashMap<String, String[]>(parameters.size() * 2);
        for (String key : parameters.keySet()) {
            map.put(key, parameters.get(key).clone());
        }
        return map;
    }

    private void parseRequest() throws IOException {
        if (parameters != null) {
            return; //already parsed
        }

        HashMap<String, List<String>> mapA = Maps.newHashMap();
        List<String> list;

        Map<String, List<String>> query = HTTPRequestUtils.getInstance().getQueryParams();
        if (query != null) {
            query.forEach((k, v) -> mapA.put(k, v));
        }

        // Fix Bug 解决Post FormData 参数找不到问题
        Map<String, String[]> parameterMap = req.getParameterMap();
        parameterMap.forEach((k, v) -> {
            if (k.contains("Content-Disposition: form-data; name")) {
                StringTokenizer st = new StringTokenizer(v[0], "\r\n");
                String key = "";
                String val = "";
                while (st.hasMoreTokens()) {
                    String token = st.nextToken();
                    if (!token.contains("------WebKitFormBoundary")) {
                        if (StringUtils.isNotBlank(key)) {
                            val = token;
                        } else {
                            key = token.replaceAll("Content-Disposition: form-data; name=", "");
                            if (key.startsWith("\"") && key.endsWith("\"")) {
                                key = key.substring(1);
                                key = key.substring(0, key.length() - 1);
                            }
                        }
                    } else {
                        if (StringUtils.isNotBlank(key) && StringUtils.isNotBlank(val)) {
                            LOG.info(" key: {} val: {}", key, val);
                            mapA.put(key, Lists.newArrayList(val));
                        }
                        key = "";
                    }
                }

            } else {
                mapA.put(k, Lists.newArrayList(v));
            }
        });

        if (shouldBufferBody()) {

            // Read the request body inputstream into a byte array.
            ByteArrayOutputStream baos = new ByteArrayOutputStream();
            try {
                // Copy all bytes from inputstream to byte array, and record time taken.
                long bufferStartTime = System.nanoTime();
                IOUtils.copy(req.getInputStream(), baos);
                bodyBufferingTimeNs = System.nanoTime() - bufferStartTime;

                contentData = baos.toByteArray();
            } catch (SocketTimeoutException e) {
                // This can happen if the request body is smaller than the size specified in the
                // Content-Length header, and using tomcat APR connector.
                LOG.error("SocketTimeoutException reading request body from inputstream. error=" + String.valueOf(e.getMessage()));
                if (contentData == null) {
                    contentData = new byte[0];
                }
            }

            try {
                LOG.debug("Length of contentData byte array = " + contentData.length);
                if (req.getContentLength() != contentData.length) {
                    LOG.warn("Content-length different from byte array length! cl=" + req.getContentLength() + ", array=" + contentData.length);
                }
            } catch (Exception e) {
                LOG.error("Error checking if request body gzipped!", e);
            }

            final boolean isPost = req.getMethod().equals("POST");

            String contentType = req.getContentType();
            final boolean isFormBody = contentType != null && contentType.contains("application/x-www-form-urlencoded");

            // only does magic body param parsing for POST form bodies
            if (isPost && isFormBody && contentData.length > 0) {
                String enc = req.getCharacterEncoding();
                if (enc == null) {
                    enc = "UTF-8";
                }
                String s = new String(contentData, enc), name, value;
                StringTokenizer st = new StringTokenizer(s, "&");
                int i;

                boolean decode = req.getContentType() != null;
                while (st.hasMoreTokens()) {
                    s = st.nextToken();
                    i = s.indexOf("=");
                    if (i > 0 && s.length() > i + 1) {
                        name = s.substring(0, i);
                        value = s.substring(i + 1);
                        if (decode) {
                            try {
                                name = URLDecoder.decode(name, "UTF-8");
                            } catch (Exception e) {
                            }
                            try {
                                value = URLDecoder.decode(value, "UTF-8");
                            } catch (Exception e) {
                            }
                        }
                        list = mapA.get(name);
                        if (list == null) {
                            list = new LinkedList<String>();
                            mapA.put(name, list);
                        }
                        list.add(value);
                    }
                }
            }
        }

        HashMap<String, String[]> map = new HashMap<String, String[]>(mapA.size() * 2);
        for (String key : mapA.keySet()) {
            list = mapA.get(key);
            map.put(key, list.toArray(new String[list.size()]));
        }

        parameters = map;

    }

    private boolean shouldBufferBody() {

        if (LOG.isDebugEnabled()) {
            LOG.debug("Path = " + req.getPathInfo());
            LOG.debug("Transfer-Encoding = " + String.valueOf(req.getHeader(ZuulHeaders.TRANSFER_ENCODING)));
            LOG.debug("Content-Encoding = " + String.valueOf(req.getHeader(ZuulHeaders.CONTENT_ENCODING)));
            LOG.debug("Content-Length header = " + req.getContentLength());
        }

        boolean should = false;
        if (req.getContentLength() > 0) {
            should = true;
        } else if (req.getContentLength() == -1) {
            final String transferEncoding = req.getHeader(ZuulHeaders.TRANSFER_ENCODING);
            if (transferEncoding != null && transferEncoding.equals(ZuulHeaders.CHUNKED)) {
                RequestContext.getCurrentContext().setChunkedRequestBody();
                should = true;
            }
        }

        return should;
    }

    /**
     * Time taken to buffer the request body in nanoseconds.
     *
     * @return
     */
    @Override
    public long getBodyBufferingTimeNs() {
        return bodyBufferingTimeNs;
    }

    /**
     * This method is safe to call multiple times.
     * Calling it will not interfere with getParameterXXX() or getReader().
     * Every time a new ServletInputStream is returned that reads data from the begining.
     *
     * @return A new ServletInputStream.
     */
    @Override
    public ServletInputStream getInputStream() throws IOException {
        parseRequest();

        return new ServletInputStreamWrapper(contentData);
    }

    /**
     * This method is safe to call multiple times.
     * Calling it will not interfere with getParameterXXX() or getInputStream().
     * Every time a new BufferedReader is returned that reads data from the begining.
     *
     * @return A new BufferedReader with the wrapped request's character encoding (or UTF-8 if null).
     */
    @Override
    public BufferedReader getReader() throws IOException {
        parseRequest();

        String enc = req.getCharacterEncoding();
        if (enc == null) {
            enc = "UTF-8";
        }
        return new BufferedReader(new InputStreamReader(new ByteArrayInputStream(contentData), enc));
    }

    /**
     * This method is safe to execute multiple times.
     *
     * @see javax.servlet.ServletRequest#getParameter(java.lang.String)
     */
    @Override
    public String getParameter(String name) {
        try {
            parseRequest();
        } catch (IOException e) {
            throw new IllegalStateException("Cannot parse the request!", e);
        }
        if (parameters == null) {
            return null;
        }
        String[] values = parameters.get(name);
        if (values == null || values.length == 0) {
            return null;
        }
        return values[0];
    }

    /**
     * This method is safe.
     *
     * @see {@link #getParameters()}
     * @see javax.servlet.ServletRequest#getParameterMap()
     */
    @SuppressWarnings("unchecked")
    @Override
    public Map getParameterMap() {
        try {
            parseRequest();
        } catch (IOException e) {
            throw new IllegalStateException("Cannot parse the request!", e);
        }
        return getParameters();
    }

    /**
     * This method is safe to execute multiple times.
     *
     * @see javax.servlet.ServletRequest#getParameterNames()
     */
    @SuppressWarnings("unchecked")
    @Override
    public Enumeration getParameterNames() {
        try {
            parseRequest();
        } catch (IOException e) {
            throw new IllegalStateException("Cannot parse the request!", e);
        }
        return new Enumeration<String>() {
            private String[] arr = getParameters().keySet().toArray(new String[0]);
            private int idx = 0;

            @Override
            public boolean hasMoreElements() {
                return idx < arr.length;
            }

            @Override
            public String nextElement() {
                return arr[idx++];
            }

        };
    }

    /**
     * This method is safe to execute multiple times.
     * Changing the returned array will not interfere with this class operation.
     *
     * @see javax.servlet.ServletRequest#getParameterValues(java.lang.String)
     */
    @Override
    public String[] getParameterValues(String name) {
        try {
            parseRequest();
        } catch (IOException e) {
            throw new IllegalStateException("Cannot parse the request!", e);
        }
        if (parameters == null) {
            return null;
        }
        String[] arr = parameters.get(name);
        if (arr == null) {
            return null;
        }
        return arr.clone();
    }
}